import os
import glob
import cv2
import numpy as np
from PIL import Image
from paddleocr import PaddleOCR


# -------------------------- 函数1：获取图片路径 --------------------------
def get_image_paths_from_folder(folder_path, supported_formats=None):
    if supported_formats is None:
        supported_formats = ['.jpg', '.jpeg', '.png', '.bmp']
    
    image_paths = []
    for filename in os.listdir(folder_path):
        file_ext = os.path.splitext(filename)[1].lower()
        if file_ext in supported_formats:
            full_path = os.path.abspath(os.path.join(folder_path, filename))
            image_paths.append(full_path)
    
    image_paths = list(set(image_paths))
    image_paths.sort()
    
    if len(image_paths) == 0:
        print(f"提示：在文件夹 {folder_path} 中未找到支持格式的图片")
    else:
        print(f"成功获取 {len(image_paths)} 张图片路径")
    
    return image_paths


# -------------------------- 函数2：OCR识别 --------------------------
def ocr_single_image(image_path, output_dir="ocr_results", use_light_model=True):
    if use_light_model:
        ocr = PaddleOCR(
            text_detection_model_name="PP-OCRv5_mobile_det",
            text_recognition_model_name="PP-OCRv5_mobile_rec",
            use_doc_orientation_classify=False,
            use_doc_unwarping=False,
            use_textline_orientation=False,
            lang="ch"
        )
    else:
        ocr = PaddleOCR(
            use_doc_orientation_classify=False,
            use_doc_unwarping=False,
            use_textline_orientation=False,
            lang="ch"
        )
    
    os.makedirs(output_dir, exist_ok=True)
    img_basename = os.path.splitext(os.path.basename(image_path))[0]
    img_output_dir = os.path.join(output_dir, img_basename)
    os.makedirs(img_output_dir, exist_ok=True)
    
    try:
        print(f"\n正在识别图片：{os.path.basename(image_path)}")
        ocr_result = ocr.predict(image_path)[0]
    except Exception as e:
        print(f"识别图片 {os.path.basename(image_path)} 失败：{str(e)}")
        return None
    
    ocr_result.save_to_img(os.path.join(img_output_dir, f"{img_basename}_annotated"))
    ocr_result.save_to_json(os.path.join(img_output_dir, f"{img_basename}_result"))
    print(f"识别结果已保存至：{img_output_dir}")
    
    structured_result = []
    for line in ocr_result:
        structured_result.append({
            "text": line[0],
            "coordinates": line[1],
            "confidence": line[2]
        })
    
    return structured_result


# -------------------------- 主逻辑（核心修改：使用英文目录保存裁剪图） --------------------------
if __name__ == "__main__":
    img_folder = "./my_images"  # 图片文件夹路径（若有中文，尽量改为英文）
    all_img_paths = get_image_paths_from_folder(folder_path=img_folder)
    
    for img_path in all_img_paths:
        img_name = os.path.basename(img_path)
        
        # 用PIL打开图片
        try:
            pil_img = Image.open(img_path)
            img_width, img_height = pil_img.size
            print(f"成功读取图片：{img_name}（尺寸：宽{img_width} × 高{img_height}）")
        except Exception as e:
            print(f"读取图片 {img_name} 失败：{str(e)}")
            continue
        
        # 计算裁剪坐标
        # crop_left = int(img_width * (1/4))    # 左侧1/4
        crop_left = int(1)    # 左侧1/4
        # crop_right = int(img_width * (3/4))   # 右侧1/4
        crop_right = int(img_width)   # 右侧1/4
        crop_upper = int(img_height - 400)    # 底部往上400像素
        crop_lower = int(img_height)          # 图片底部
        crop_coords = (crop_left, crop_upper, crop_right, crop_lower)
        
        # 验证坐标合法性
        valid = True
        if crop_left < 0 or crop_upper < 0 or crop_right > img_width or crop_lower > img_height:
            valid = False
            print(f"警告：裁剪坐标超出图片范围！坐标：{crop_coords}，图片尺寸：(宽{img_width}, 高{img_height})")
        if crop_left >= crop_right or crop_upper >= crop_lower:
            valid = False
            print(f"警告：裁剪坐标逻辑错误！left={crop_left} ≥ right={crop_right} 或 upper={crop_upper} ≥ lower={crop_lower}")
        
        # 裁剪或使用原图（核心修改：用英文目录保存裁剪图）
        if valid:
            cropped_pil_img = pil_img.crop(crop_coords)
            # 改为英文目录，避免中文路径问题
            crop_dir = "./temp_crop"  # 关键：删除中文，用英文目录
            os.makedirs(crop_dir, exist_ok=True)
            # 裁剪后的文件名也避免中文（若原文件名有中文，可保留，但路径必须英文）
            cropped_img_path = os.path.join(crop_dir, f"cropped_{img_name}")
            cropped_pil_img.save(cropped_img_path)
            print(f"已裁剪图片，保存路径：{cropped_img_path}（裁剪区域：{crop_coords}）")
        else:
            cropped_img_path = img_path
            print(f"裁剪坐标不合法，将直接识别原图：{img_name}")
        
        # 执行OCR
        result = ocr_single_image(image_path=cropped_img_path)
        
        # 输出结果
        if result:
            all_text = [item["text"] for item in result]
            print(f"图片 {os.path.basename(cropped_img_path)} 识别到的文字：{all_text}")

    print("\n所有图片处理完成！")
